""" Upload data such as samples to MISP."""
import os
import magic
import zipfile
import json
import pymisp
import traceback
import ssl
import sys
from datetime import datetime
from requests.exceptions import ConnectionError
from urllib3.exceptions import NewConnectionError
from pymisp.exceptions import MISPServerError
from pathlib import Path
from mispconn.utils import logger, CFG, parse_excluding
from tools import Tools
from pymisp.tools import make_binary_objects, FileObject
from typing import TypeVar, Tuple, List, Dict, Union, Generator, Sequence, Set, \
    Optional
from mispconn.apk_parse3.apk import APK
FILE_TYPE_MAP: dict = dict()
KNOWN_FILE_TYPES: Set[str] = set()
BIN_TYPES = {"peexe", "peexe32", "pedll", "peexe64", "elf", "elf64", "elf32"}
TAG_NAMESPACE: str = CFG['misp']['tag_namespace']
ssl._create_default_https_context = ssl._create_unverified_context


class SampleUploader:
    def __init__(self, path: str, tags: Union[Sequence[str], str, None],
                 delete: bool, file_type: Union[Sequence[str], str],
                 file_paths_for_excluding: Union[Sequence[str], str, None]):
        """
        Upload malware samples along with tags to misp. Create one event for
        each sample. Attributes are populated according to the file type.
        :param path: where the malware samples are.
        :param tags: tags to event and malware-sample attribute.
        :param delete: whether to delete the origin sample file after upload.
        :param file_type: file type of sample such as elf, peexe.
        :param file_paths_for_excluding: the path of text file including paths
                to be excluded. This will speed up tm.
        """
        self.root_path = path
        self.to_delete = delete
        self.tags = tuple(tags) if tags else tuple()
        self.wanted_file_types = file_type \
            if isinstance(file_type, list) else [file_type]
        self.file_path_for_future_excluding: Optional[str] = None
        self.excluded_paths: Set[str] = set()
        self.init_excluding(file_paths_for_excluding)
        self.misp = Tools.init_misp(url=CFG['misp']['url'],
                                    key=CFG['misp']['key'])
        self.load_file_types()
        self.execute()

    def init_excluding(self, file_paths_for_excluding:
    Union[Sequence[str], str, None]):
        # Create a new file to record processed paths for this run, which
        # can be used in the future runs.
        fname = 'aicti_tm_excluded_paths_in' + \
                self.root_path.replace('/', '_') + '_' + \
                datetime.now().strftime('%Y%m%d%H%M%S')
        self.file_path_for_future_excluding = os.path.join('/tmp', fname)
        logger.info(f'processed paths will be recorded in '
                    f'{self.file_path_for_future_excluding} for future '
                    f'exclusion.')
        self.excluded_paths = parse_excluding(file_paths_for_excluding)

    def execute(self):
        for file_path in self.get_file_paths():
            try:
                # Ignore unwanted files
                if file_path in self.excluded_paths:
                    continue
                if not self.is_valid_size(file_path):
                    continue
                if not self.is_wanted_file_type(file_path):
                    continue
                if self.is_in_existence(file_path):
                    self.delete_file(file_path)
                    continue
                self.upload(file_path)
                self.delete_file(file_path)
            except (NewConnectionError, ConnectionError) as e:
                logger.critical(f'{e}')
                sys.exit('ConnectionError')
            except MISPServerError as e:
                if 'Error code 500' in str(e):
                    logger.critical(f'{e}')
                    sys.exit('ConnectionError')
                else:
                    logger.error(f'{e}, {file_path}')
            except Exception as e:
                logger.error(f'{e}, {file_path}')
            finally:
                self.add_to_file_for_excluding(file_path)
        logger.info('finished, all files are uploaded.')

    def delete_file(self, file_path):
        if self.to_delete:
            os.remove(file_path)
            logger.debug(f'removed {file_path}')

    def add_to_file_for_excluding(self, file_path: str):
        with open(self.file_path_for_future_excluding, 'a') as f:
            f.write(f'{file_path}\n')

    def is_wanted_file_type(self, file_path: str) -> bool:
        # If wanted file types are known
        if 'all' in self.wanted_file_types:
            return True
        if not set(self.wanted_file_types).issubset(KNOWN_FILE_TYPES):
            msg = f'Unknown file type {self.wanted_file_types}!\n' \
                  f'Valid types are: {KNOWN_FILE_TYPES}'
            sys.exit(msg)
        # If this file is the wanted type
        file_type, _ = self.determine_file_type_processor(file_path)
        if bool(set.intersection(set(file_type), set(self.wanted_file_types))):
            return True
        else:
            return False

    @staticmethod
    def is_valid_size(file_path: str) -> bool:
        """Ignore empty and too large files"""
        result = True
        file_size = os.path.getsize(file_path)
        if file_size == 0:
            result = False
        elif file_size > CFG['misp'].getint('max_sample_size'):
            result = False
            logger.info(f'size of {file_path} is over the limit %s'
                        % (CFG['misp'].getint('max_sample_size')))
        return result

    def is_in_existence(self, file_path: str) -> bool:
        md5sum = Tools.md5sum(file_path)
        res = self.misp.search(controller='objects',
                               type_attribute='malware-sample',
                               value=md5sum)
        if len(res) > 0:
            logger.debug(f'{file_path} was already in misp, ignores.')
            return True
        else:
            return False

    @staticmethod
    def load_file_types():
        # Load mapping from magic file type to file type tag in MISP.
        global FILE_TYPE_MAP, KNOWN_FILE_TYPES
        __location__ = os.path.realpath(
            os.path.join(os.getcwd(), os.path.dirname(__file__)))
        with open(os.path.join(__location__, 'file_type_map.json')) as f:
            FILE_TYPE_MAP = json.load(f)
        KNOWN_FILE_TYPES = set()
        for v in FILE_TYPE_MAP.values():
            if isinstance(v, list):
                KNOWN_FILE_TYPES = set.union(KNOWN_FILE_TYPES, set(list(v)))
            elif isinstance(v, str):
                KNOWN_FILE_TYPES.add(v)
            else:
                raise RuntimeError('Incorrect file type')

    def get_file_paths(self) -> Generator[str, None, None]:
        """yield paths of all files in the given path"""
        for root, dirs, files in os.walk(self.root_path):
            for file in files:
                yield os.path.join(root, file)

    @staticmethod
    def determine_file_type_processor(file_path: str) -> \
            (Union[List[str], str, None], Optional[str]):
        """Determine the file type and add it to tags."""
        # file_type = magic.from_file(file_path).strip().split(',')[0].strip()
        file_type, processor = None, None
        magic_res = magic.from_file(file_path)
        assert type(magic_res) is str and len(magic_res) > 0
        origin_file_type = magic_res.split(',')[0].strip()
        file_type = origin_file_type
        if origin_file_type.startswith('Zip') or \
                origin_file_type.startswith('Java archive data (JAR)'):
            # Determine apk
            try:
                with zipfile.ZipFile(file_path, 'r') as zf:
                    zf.getinfo('AndroidManifest.xml')
                    origin_file_type = 'APK'
            except Exception as e:
                pass
        # Mapping from magic file type to file type tag in MISP.
        for k, v in FILE_TYPE_MAP.items():
            if origin_file_type.startswith(k):
                file_type = v
                break
        if isinstance(file_type, str):
            file_type = [file_type]
        # Determine processor of elf files. PE32 and PE32+ are ignored
        # because they are default to Intel 80386 and x86-64 respectively.
        if origin_file_type.startswith('ELF'):
            magic_res_split = magic_res.split(',')
            if len(magic_res_split) > 0:
                processor = magic_res_split[1].strip()
        return file_type, processor

    def add_binary_objects(self, file_path: str, event: pymisp.MISPEvent):
        # Populate objects for binary samples
        fo, peo, seos = make_binary_objects(file_path)
        if seos:
            for s in seos:
                r = self.misp.add_object(event, s)
        if peo:
            r = self.misp.add_object(event, peo, pythonify=True)
            for ref in peo.ObjectReference:
                r = self.misp.add_object_reference(ref)
        if fo:
            # response = self.misp.add_object(event, fo, pythonify=True)
            for ref in fo.ObjectReference:
                r = self.misp.add_object_reference(ref)

    def add_apk_attributes(self, file_path: str, event: pymisp.MISPEvent):
        # Populate attributes of APK files
        # Change default category of attribute type malware-sample
        describe_types = self.misp.describe_types_local
        describe_types['sane_defaults']['text']['default_category'] \
            = 'External analysis'
        try:
            attrs = dict()
            apkf = APK(file_path)
            attrs['app-name'] = apkf.get_app_name()
            attrs['package-name'] = apkf.get_package()
            try:
                attrs['cert-sha1'] = apkf.get_cert_sha1()[0]
            except:
               pass
            for k, v in attrs.items():
                attr = pymisp.MISPAttribute(describe_types=describe_types)
                attr.type = 'threat-actor'
                attr.value = f'{k}="{v}"'
                self.misp.add_attribute(event, attr)

        except Exception as e:
            logger.error(f'failed to parse {file_path}')

    def upload(self, file_path: str):
        """Upload file to misp."""
        event = pymisp.MISPEvent()
        event.info = f'{Path(file_path).name}'
        event = self.misp.add_event(event)
        event_uuid = event['Event']['uuid']
        # Generate file object
        fo = self.misp.add_object(event, FileObject(filepath=(Path(file_path))))
        file_type, processor = self.determine_file_type_processor(file_path)
        # Generate binary objects
        if bool(set.intersection(BIN_TYPES, set(file_type))) \
                and CFG['misp']['make_binary_objects'] == 'True':
            self.add_binary_objects(file_path, event)
        # Generate APK attributes
        if 'android' in file_type:
            self.add_apk_attributes(file_path, event)
        # Tag sample
        attr_sample_uuid = ''
        for attr in fo['Object']['Attribute']:
            if attr['type'] == 'malware-sample':
                attr_sample_uuid = attr['uuid']
                break
        assert len(attr_sample_uuid) == len(event_uuid) == 36
        tags = list(self.tags)
        if processor is not None:
            tags.append(f'{TAG_NAMESPACE}:processor="{processor}"')
        if set(file_type).issubset(KNOWN_FILE_TYPES):
            for t in file_type:
                tags.append(f'file-type:type="{t}"')
        else:
            for t in file_type:
                tags.append(f'{TAG_NAMESPACE}:file-type="{t}"')

        for tag in tags:
            self.misp.tag(attr_sample_uuid, tag)
            self.misp.tag(event_uuid, tag)
        logger.debug(f'uploaded {file_path} to misp successfully.')